import argparse
import json
import base64
import os
import secrets
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

PRIVATE_KEY_PATH = "private_key.pem"
PUBLIC_KEY_PATH = "public_key.pem"
CERTIFICATE_PATH = "license.json"
ENCRYPTED_CERTIFICATE_PATH = "encrypted_license.json"
CLIENT_SEED_PATH = "client_decrypt_seed.json"
# 方案C相关文件路径
RSA_ENCRYPTED_PATH = "encrypted.txt"

# 双种子配置（方案B）
# 注意：这里实际上使用相同的种子，但名称不同是为了概念上的区分
ENCRYPTION_SEED = "license-encryption-seed-2024"  # 用于加密和解密的种子
CLIENT_DECRYPT_SEED = "license-encryption-seed-2024"  # 给客户端的解密种子（实际相同）


def generate_and_save_keys():
    """生成并保存RSA密钥对"""
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    public_key = private_key.public_key()

    with open(PRIVATE_KEY_PATH, "wb") as f:
        f.write(
            private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption(),
            )
        )

    with open(PUBLIC_KEY_PATH, "wb") as f:
        f.write(
            public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo,
            )
        )

    print(f"✅ 密钥对已重新生成：")
    print(f"   - 私钥：{PRIVATE_KEY_PATH}")
    print(f"   - 公钥：{PUBLIC_KEY_PATH}")
    return private_key, public_key


def load_private_key():
    """加载私钥"""
    with open(PRIVATE_KEY_PATH, "rb") as f:
        return serialization.load_pem_private_key(f.read(), password=None)


def load_public_key():
    """加载公钥"""
    with open(PUBLIC_KEY_PATH, "rb") as f:
        return serialization.load_pem_public_key(f.read())


def derive_encryption_key(user_info, salt, seed):
    """从用户信息派生AES密钥（使用指定种子）"""
    password = f"{user_info}#{seed}".encode("utf-8")

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,  # 256-bit key
        salt=salt,
        iterations=100000,
    )

    return kdf.derive(password)


def encrypt_with_aes(data, key, iv):
    """使用AES-256-GCM加密数据"""
    cipher = Cipher(algorithms.AES(key), modes.GCM(iv))
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(data.encode("utf-8")) + encryptor.finalize()
    return ciphertext, encryptor.tag


def decrypt_with_aes(ciphertext, key, iv, tag):
    """使用AES-256-GCM解密数据"""
    cipher = Cipher(algorithms.AES(key), modes.GCM(iv, tag))
    decryptor = cipher.decryptor()
    plaintext = decryptor.update(ciphertext) + decryptor.finalize()
    return plaintext.decode("utf-8")


def sign_certificate(json_string, output_file=CERTIFICATE_PATH):
    """功能1：生成公钥私钥和数字证书（每次重新生成三个文件）"""
    print("🔑 正在重新生成密钥对和数字证书...")

    # 每次都重新生成密钥对
    private_key, public_key = generate_and_save_keys()

    try:
        # 对许可证数据进行数字签名
        signature = private_key.sign(
            json_string.encode("utf-8"),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256(),
        )

        # 将签名转换为 Base64
        signature_b64 = base64.b64encode(signature).decode("utf-8")

        # 创建数字证书（方案A：明文数据 + 数字签名）
        certificate = {
            "data": json_string,  # 明文许可证数据
            "signature": signature_b64,
            "algorithm": "RSA-PSS-SHA256",
            "version": "1.0",
            "type": "license_certificate",
        }

        # 保存数字证书
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(certificate, f, ensure_ascii=False, indent=2)

        print(f"✅ 数字证书生成成功：")
        print(f"   - 数字证书：{output_file}")
        print(f"📝 签名算法：RSA-PSS-SHA256")

        print(f"\n🎯 生成的文件清单：")
        print(f"   1. {PRIVATE_KEY_PATH} - 私钥（您保留）")
        print(f"   2. {PUBLIC_KEY_PATH} - 公钥（给客户端）")
        print(f"   3. {output_file} - 数字证书（给客户端）")

        return True

    except Exception as e:
        print(f"❌ 数字证书生成失败: {e}")
        return False


def sign_encrypted_certificate(json_string, output_file=ENCRYPTED_CERTIFICATE_PATH):
    """方案B：生成加密数字证书（data字段AES加密）"""
    print("🔑 正在重新生成密钥对和加密数字证书...")

    # 每次都重新生成密钥对
    private_key, public_key = generate_and_save_keys()

    try:
        # 解析用户信息用于密钥派生
        license_data = json.loads(json_string)
        user_info = license_data.get("user", "default")

        # 1. 对原始明文数据进行签名
        signature = private_key.sign(
            json_string.encode("utf-8"),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256(),
        )

        # 2. 生成随机盐值和IV
        salt = secrets.token_bytes(16)
        iv = secrets.token_bytes(12)  # GCM mode uses 12-byte IV

        # 3. 使用服务端种子派生加密密钥
        encryption_key = derive_encryption_key(user_info, salt, ENCRYPTION_SEED)

        # 4. 使用AES-GCM加密data字段
        ciphertext, tag = encrypt_with_aes(json_string, encryption_key, iv)

        # 转换为Base64
        signature_b64 = base64.b64encode(signature).decode("utf-8")
        ciphertext_b64 = base64.b64encode(ciphertext).decode("utf-8")
        salt_b64 = base64.b64encode(salt).decode("utf-8")
        iv_b64 = base64.b64encode(iv).decode("utf-8")
        tag_b64 = base64.b64encode(tag).decode("utf-8")

        # 创建加密数字证书
        encrypted_certificate = {
            "data": ciphertext_b64,  # AES加密的数据
            "signature": signature_b64,
            "algorithm": "RSA-PSS-SHA256",
            "encryption": "AES-256-GCM",
            "salt": salt_b64,
            "iv": iv_b64,
            "tag": tag_b64,
            "user_hint": user_info,
            "version": "1.0",
            "type": "encrypted_license_certificate",
        }

        # 保存加密数字证书
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(encrypted_certificate, f, ensure_ascii=False, indent=2)

        print(f"✅ 加密数字证书生成成功：")
        print(f"   - 加密证书：{output_file}")
        print(f"📝 签名算法：RSA-PSS-SHA256")
        print(f"🔒 加密算法：AES-256-GCM")
        print(f"🔑 用户标识：{user_info}")

        # 生成客户端解密种子文件（只包含解密种子）
        client_seed_data = {
            "decrypt_seed": CLIENT_DECRYPT_SEED,
            "description": "客户端解密专用种子（不能用于加密）",
            "algorithm": "PBKDF2-SHA256",
            "security_note": "此种子仅用于解密，无法用于伪造许可证",
        }

        with open(CLIENT_SEED_PATH, "w", encoding="utf-8") as f:
            json.dump(client_seed_data, f, ensure_ascii=False, indent=2)

        print(f"\n🎯 生成的文件清单：")
        print(f"   1. {PRIVATE_KEY_PATH} - 私钥（您保留）")
        print(f"   2. {PUBLIC_KEY_PATH} - 公钥（给客户端）")
        print(f"   3. {output_file} - 加密证书（给客户端）")
        print(f"   4. {CLIENT_SEED_PATH} - 解密种子（给客户端）")

        print(f"\n🛡️  方案B安全特性：")
        print(f"   ✅ 数据完全加密保护")
        print(f"   ✅ 客户端只能解密，无法伪造")
        print(f"   ✅ 双种子安全机制")

        return True

    except Exception as e:
        print(f"❌ 加密数字证书生成失败: {e}")
        return False


def verify_certificate(cert_file=CERTIFICATE_PATH):
    """功能2：验证数字证书"""
    print("🔍 正在验证数字证书...")

    if not os.path.exists(PUBLIC_KEY_PATH):
        print(f"❌ 公钥文件 {PUBLIC_KEY_PATH} 不存在")
        print("💡 提示：请先运行生成功能，或确保公钥文件存在")
        return False

    if not os.path.exists(cert_file):
        print(f"❌ 数字证书文件 {cert_file} 不存在")
        return False

    try:
        # 读取数字证书
        with open(cert_file, "r", encoding="utf-8") as f:
            certificate = json.load(f)

        # 提取证书字段
        data = certificate.get("data")
        signature_b64 = certificate.get("signature")
        algorithm = certificate.get("algorithm", "RSA-PSS-SHA256")
        cert_type = certificate.get("type", "unknown")
        version = certificate.get("version", "1.0")

        if not data or not signature_b64:
            print("❌ 数字证书格式错误，缺少必要字段")
            return False

        print(f"📄 证书信息：")
        print(f"   类型：{cert_type}")
        print(f"   版本：{version}")
        print(f"   算法：{algorithm}")

        # 验证数字签名
        signature = base64.b64decode(signature_b64)
        public_key = load_public_key()

        public_key.verify(
            signature,
            data.encode("utf-8"),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256(),
        )

        print(f"✅ 数字证书验证成功！")
        print(f"🛡️  数字签名有效，数据未被篡改")

        # 显示许可证详情
        try:
            license_data = json.loads(data)
            print(f"📋 许可证详情：")
            for key, value in license_data.items():
                print(f"   {key}: {value}")
        except json.JSONDecodeError:
            print(f"📋 许可证数据：{data}")

        return True

    except FileNotFoundError:
        print(f"❌ 证书文件 {cert_file} 不存在")
        return False
    except Exception as e:
        print(f"❌ 数字证书验证失败: {e}")
        print("🚨 可能原因：")
        print("   - 证书已被篡改")
        print("   - 公钥与证书不匹配")
        print("   - 证书格式损坏")
        return False


def verify_encrypted_certificate(
    cert_file=ENCRYPTED_CERTIFICATE_PATH, seed_file=CLIENT_SEED_PATH
):
    """方案B：验证加密数字证书"""
    print("🔍 正在验证加密数字证书...")

    if not os.path.exists(PUBLIC_KEY_PATH):
        print(f"❌ 公钥文件 {PUBLIC_KEY_PATH} 不存在")
        return False

    if not os.path.exists(cert_file):
        print(f"❌ 加密证书文件 {cert_file} 不存在")
        return False

    if not os.path.exists(seed_file):
        print(f"❌ 解密种子文件 {seed_file} 不存在")
        return False

    try:
        # 读取加密数字证书
        with open(cert_file, "r", encoding="utf-8") as f:
            certificate = json.load(f)

        # 读取客户端解密种子
        with open(seed_file, "r", encoding="utf-8") as f:
            seed_data = json.load(f)

        # 提取证书字段
        ciphertext_b64 = certificate.get("data")
        signature_b64 = certificate.get("signature")
        algorithm = certificate.get("algorithm", "RSA-PSS-SHA256")
        encryption = certificate.get("encryption", "AES-256-GCM")
        salt_b64 = certificate.get("salt")
        iv_b64 = certificate.get("iv")
        tag_b64 = certificate.get("tag")
        user_hint = certificate.get("user_hint", "default")
        cert_type = certificate.get("type", "unknown")
        version = certificate.get("version", "1.0")

        if not all([ciphertext_b64, signature_b64, salt_b64, iv_b64, tag_b64]):
            print("❌ 加密证书格式错误，缺少必要字段")
            return False

        print(f"📄 证书信息：")
        print(f"   类型：{cert_type}")
        print(f"   版本：{version}")
        print(f"   签名算法：{algorithm}")
        print(f"   加密算法：{encryption}")
        print(f"   用户标识：{user_hint}")

        # 1. 使用客户端解密种子派生密钥
        decrypt_seed = seed_data.get("decrypt_seed")
        salt = base64.b64decode(salt_b64)
        decryption_key = derive_encryption_key(user_hint, salt, decrypt_seed)

        # 2. 解密data字段
        ciphertext = base64.b64decode(ciphertext_b64)
        iv = base64.b64decode(iv_b64)
        tag = base64.b64decode(tag_b64)

        try:
            original_data = decrypt_with_aes(ciphertext, decryption_key, iv, tag)
            print(f"🔓 数据解密成功")
        except Exception as decrypt_error:
            print(f"❌ 数据解密失败: {decrypt_error}")
            print("🚨 可能原因：")
            print("   - 解密种子不匹配")
            print("   - 加密数据已损坏")
            print("   - 用户标识不正确")
            return False

        # 3. 验证签名（使用解密后的原始数据）
        try:
            signature = base64.b64decode(signature_b64)
            public_key = load_public_key()

            public_key.verify(
                signature,
                original_data.encode("utf-8"),
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH,
                ),
                hashes.SHA256(),
            )
            print(f"✅ 数字签名验证成功！")
        except Exception as verify_error:
            print(f"❌ 数字签名验证失败: {verify_error}")
            print("🚨 可能原因：")
            print("   - 证书已被篡改")
            print("   - 公钥与证书不匹配")
            return False

        print(f"🛡️  数字签名有效，数据未被篡改")

        # 显示许可证详情
        try:
            license_data = json.loads(original_data)
            print(f"📋 许可证详情：")
            for key, value in license_data.items():
                print(f"   {key}: {value}")
        except json.JSONDecodeError:
            print(f"📋 许可证数据：{original_data}")

        print(f"🔒 方案B特点：数据加密保护 + 双种子安全机制")
        return True

    except FileNotFoundError as e:
        print(f"❌ 文件不存在: {e}")
        return False
    except json.JSONDecodeError as e:
        print(f"❌ JSON格式错误: {e}")
        return False
    except Exception as e:
        print(f"❌ 加密证书验证失败: {e}")
        print("💡 提示: 客户端无法伪造许可证，因为使用了不同的加密种子")
        return False


def rsa_encrypt_data(json_string, output_file=RSA_ENCRYPTED_PATH):
    """方案C：使用RSA直接加密数据（每次重新生成密钥对）"""
    print("🔑 正在生成新的密钥对...")

    # 每次加密时都重新生成密钥对
    private_key, public_key = generate_and_save_keys()

    try:
        # 验证输入是否为有效JSON
        json.loads(json_string)

        # 检查数据长度限制（RSA-2048最多加密245字节）
        if len(json_string.encode("utf-8")) > 245:
            print("❌ 数据过长，RSA-2048 最多只能加密 245 字节")
            print(f"📏 当前数据长度：{len(json_string.encode('utf-8'))} 字节")
            print("💡 建议：使用方案B（AES+RSA）处理大数据")
            return False

        # 使用RSA-OAEP加密
        encrypted = public_key.encrypt(
            json_string.encode("utf-8"),
            padding.OAEP(
                mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None
            ),
        )

        # 转换为Base64并保存
        encrypted_b64 = base64.b64encode(encrypted).decode("utf-8")

        with open(output_file, "w") as f:
            f.write(encrypted_b64)

        print(f"🔐 RSA加密成功，并保存到 {output_file} 文件")
        print(f"📝 加密算法：RSA-OAEP-SHA256")
        print(f"📏 数据长度：{len(json_string.encode('utf-8'))} 字节")

        print(f"\n🎯 生成的文件清单：")
        print(f"   1. {PRIVATE_KEY_PATH} - 私钥（解密时需要）")
        print(f"   2. {PUBLIC_KEY_PATH} - 公钥（加密时使用）")
        print(f"   3. {output_file} - RSA加密文件（密文）")

        print(f"\n⚠️  方案C特点：")
        print(f"   ✅ 数据完全加密保护")
        print(f"   ❌ 无完整性验证（无法检测篡改）")
        print(f"   ❌ 数据长度限制（最多245字节）")
        print(f"   ❌ 加密性能较差")

        return True

    except json.JSONDecodeError:
        print("❌ 请输入合法的 JSON 字符串")
        return False
    except Exception as e:
        print(f"❌ RSA加密失败: {e}")
        return False


def rsa_decrypt_data(encrypted_file=RSA_ENCRYPTED_PATH):
    """方案C：使用RSA解密数据"""
    print("🔍 正在使用RSA解密数据...")

    if not os.path.exists(PRIVATE_KEY_PATH):
        print(f"❌ 私钥文件 {PRIVATE_KEY_PATH} 不存在")
        return False

    if not os.path.exists(encrypted_file):
        print(f"❌ 加密文件 {encrypted_file} 不存在")
        return False

    try:
        # 读取加密数据
        with open(encrypted_file, "r") as f:
            encrypted_b64 = f.read().strip()

        # Base64解码
        encrypted_bytes = base64.b64decode(encrypted_b64)

        # 加载私钥
        private_key = load_private_key()

        # RSA-OAEP解密
        decrypted = private_key.decrypt(
            encrypted_bytes,
            padding.OAEP(
                mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None
            ),
        )

        decrypted_text = decrypted.decode("utf-8")

        print(f"✅ RSA解密成功！")
        print(f"📝 解密算法：RSA-OAEP-SHA256")

        # 尝试格式化JSON输出
        try:
            license_data = json.loads(decrypted_text)
            print(f"📋 解密后的数据：")
            for key, value in license_data.items():
                print(f"   {key}: {value}")
        except json.JSONDecodeError:
            print(f"📋 解密后的数据：{decrypted_text}")

        print(f"\n⚠️  方案C注意：")
        print(f"   ❌ 此方案无法验证数据完整性")
        print(f"   ❌ 无法检测数据是否被篡改")
        print(f"   💡 建议使用方案A或B获得更好的安全性")

        return True

    except Exception as e:
        print(f"❌ RSA解密失败: {e}")
        print("🚨 可能原因：")
        print("   - 私钥与加密数据不匹配")
        print("   - 加密数据已损坏")
        print("   - Base64编码错误")
        return False


def rsa_decrypt_from_base64(base64_string):
    """方案C：直接从Base64字符串解密数据"""
    print("🔍 正在从Base64字符串解密数据...")

    if not os.path.exists(PRIVATE_KEY_PATH):
        print(f"❌ 私钥文件 {PRIVATE_KEY_PATH} 不存在")
        return False

    try:
        # Base64解码
        encrypted_bytes = base64.b64decode(base64_string)

        # 加载私钥
        private_key = load_private_key()

        # RSA-OAEP解密
        decrypted = private_key.decrypt(
            encrypted_bytes,
            padding.OAEP(
                mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None
            ),
        )

        decrypted_text = decrypted.decode("utf-8")

        print(f"✅ RSA解密成功！")
        print(f"📝 解密算法：RSA-OAEP-SHA256")

        # 尝试格式化JSON输出
        try:
            license_data = json.loads(decrypted_text)
            print(f"📋 解密后的数据：")
            for key, value in license_data.items():
                print(f"   {key}: {value}")
        except json.JSONDecodeError:
            print(f"📋 解密后的数据：{decrypted_text}")

        return True

    except Exception as e:
        print(f"❌ RSA解密失败: {e}")
        return False


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="数字证书生成和验证工具 (支持方案A、方案B和方案C)"
    )
    parser.add_argument("--sign", type=str, help="方案A：生成明文数字证书")
    parser.add_argument("--sign-encrypted", type=str, help="方案B：生成加密数字证书")
    parser.add_argument(
        "--verify",
        type=str,
        nargs="?",
        const=CERTIFICATE_PATH,
        help=f"验证明文数字证书（默认：{CERTIFICATE_PATH}）",
    )
    parser.add_argument(
        "--verify-encrypted",
        type=str,
        nargs="?",
        const=ENCRYPTED_CERTIFICATE_PATH,
        help=f"验证加密数字证书（默认：{ENCRYPTED_CERTIFICATE_PATH}）",
    )
    parser.add_argument(
        "--rsa-encrypt", type=str, help="方案C：使用RSA直接加密JSON数据"
    )
    parser.add_argument(
        "--rsa-decrypt",
        type=str,
        nargs="?",
        const=RSA_ENCRYPTED_PATH,
        help=f"方案C：RSA解密文件（默认：{RSA_ENCRYPTED_PATH}）",
    )
    parser.add_argument(
        "--rsa-decrypt-base64", type=str, help="方案C：直接解密Base64字符串"
    )

    args = parser.parse_args()

    if args.sign:
        try:
            # 验证输入是否为有效JSON
            json.loads(args.sign)
            sign_certificate(args.sign)
        except json.JSONDecodeError:
            print("❌ 请输入合法的 JSON 字符串")
            print('💡 示例：--sign \'{"user":"客户A","licensetype":"premium"}\'')
    elif args.sign_encrypted:
        try:
            # 验证输入是否为有效JSON
            json.loads(args.sign_encrypted)
            sign_encrypted_certificate(args.sign_encrypted)
        except json.JSONDecodeError:
            print("❌ 请输入合法的 JSON 字符串")
            print(
                '💡 示例：--sign-encrypted \'{"user":"客户A","licensetype":"premium"}\''
            )
    elif args.verify:
        verify_certificate(args.verify)
    elif args.verify_encrypted:
        verify_encrypted_certificate(args.verify_encrypted)
    elif args.rsa_encrypt:
        try:
            # 验证输入是否为有效JSON
            json.loads(args.rsa_encrypt)
            rsa_encrypt_data(args.rsa_encrypt)
        except json.JSONDecodeError:
            print("❌ 请输入合法的 JSON 字符串")
            print('💡 示例：--rsa-encrypt \'{"user":"客户A","licensetype":"premium"}\'')
    elif args.rsa_decrypt:
        try:
            # 验证输入是否为有效文件路径
            os.path.exists(args.rsa_decrypt)
            rsa_decrypt_data(args.rsa_decrypt)
        except FileNotFoundError:
            print(f"❌ 文件不存在: {args.rsa_decrypt}")
    elif args.rsa_decrypt_base64:
        try:
            # 验证输入是否为有效Base64字符串
            rsa_decrypt_from_base64(args.rsa_decrypt_base64)
        except Exception as e:
            print(f"❌ 输入错误: {e}")
    else:
        print("📖 使用说明：")
        print(
            f'  方案A - 生成明文证书：python3 {os.path.basename(__file__)} --sign \'{{"user":"客户A","licensetype":"premium"}}\''
        )
        print(f"  方案A - 验证明文证书：python3 {os.path.basename(__file__)} --verify")
        print(
            f'  方案B - 生成加密证书：python3 {os.path.basename(__file__)} --sign-encrypted \'{{"user":"客户A","licensetype":"premium"}}\''
        )
        print(
            f"  方案B - 验证加密证书：python3 {os.path.basename(__file__)} --verify-encrypted"
        )
        print(
            f"  方案C - 使用RSA加密数据：python3 {os.path.basename(__file__)} --rsa-encrypt "
            + '\'{"user":"客户A","licensetype":"premium"}\''
        )
        print(
            f"  方案C - 使用RSA解密文件：python3 {os.path.basename(__file__)} --rsa-decrypt {RSA_ENCRYPTED_PATH}"
        )
        print(
            f"  方案C - 使用RSA解密Base64字符串：python3 {os.path.basename(__file__)} --rsa-decrypt-base64 'Base64EncodedString'"
        )
        parser.print_help()
